#encoding:gbk
'''
ETF排序买入策略：
超短强度:=ATAN((MA(C,5)/REF(MA(C,5),1)-1)*100)*57.3;
短期强度:=ATAN((MA(C,13)/REF(MA(C,13),1)-1)*100)*57.3;
中期强度:=ATAN((MA(C,20)/REF(MA(C,20),1)-1)*100)*57.3;
长期强度:=ATAN((MA(C,60)/REF(MA(C,60),1)-1)*100)*57.3;
强度评分:超短强度/50+短期强度/40+中期强度/21+长期强度/10;
'''
import pandas as pd
import numpy as np
import talib
import math
import pickle
from datetime import datetime

#---- 参数设置 ----#
etf_file_path = 'C:\\new_tdx\\T0002\\blocknew\\ETF.blk'
stock_poor = {}
买入金额 = 2000
买入下单溢价 = 1.01  # 价格上浮1%
卖出下单溢价 = 0.99  # 价格下浮1%
now = datetime.now()
now = now.strftime('%Y%m%d')
etf_bak_file = 'C:\\logs\\etf_bak_file.plk'

def init(ContextInfo):
    # ContextInfo.trade_code_list=['601398.SH','601857.SH','601288.SH','000333.SZ','002415.SZ','000002.SZ']
    # ContextInfo.set_universe(ContextInfo.trade_code_list)
    ContextInfo.accID = 'XXXXXXXX'
    ContextInfo.start = '2024-11-01 00:00:00'
    ContextInfo.end = '2024-11-01 15:00:00'
    ContextInfo.run_time("trade","1nDay","2024-10-25 14:50:00")

# ETF排序，返回排序后前十的列表
def etf_sort(ContextInfo, etf_file_path):
    with open(etf_file_path, 'r') as f:
        data = f.readlines()
    # data = ContextInfo.get_stock_list_in_sector('ETF池')
    for line in data:
        if len(line) < 6:
            continue
        line = line.replace('\n', '')
        if line[0] == '0':
            # stock_poor.append(line[1:] + '.SZ')
            stock_poor[line[1:] + '.SZ'] ={}
        if line[0] == '1':
            # stock_poor.append(line[1:] + '.SH')
            stock_poor[line[1:] + '.SH'] ={}
        if line[0] == '2':
            # stock_poor.append(line[1:] + '.BJ')
            stock_poor[line[1:] + '.BJ'] ={}
    
    stk_list = list(stock_poor.keys())
    # stk_list = ['600785.SH','601500.SH','510900.SH','588200.SH','513300.SH','600171.SH','000777.SZ','002555.SZ','002558.SZ','300059.SZ','600064.SH','001872.SZ','301012.SZ','002588.SZ','603055.SH','600056.SH','000705.SZ','000923.SZ']
    # for stk in stk_list:
        # stock_poor[stk] = {}
    # print(stk_list)
    # print(type(stk_list))
    data = ContextInfo.get_market_data_ex(
        fields = ['close'],
        stock_code = stk_list,
        period = '1d',
        # start_time = '',
        end_time = '',
        count = 62,
        dividend_type = 'follow',
        fill_data = True,
        subscribe = False)
    
    p_tmp = {}
    ticks = ContextInfo.get_full_tick(stk_list)
    for stk in stk_list:
        # stock_poor[stk]['ma5'] = data[stk]['close'][-5:].mean()
        stock_poor[stk]['ma5'] = (data[stk]['close'][-5:-1].sum() + ticks[stk]['lastPrice'])/5
        stock_poor[stk]['ma5-1'] = data[stk]['close'][-6:-1].mean()
        # stock_poor[stk]['ma13'] = data[stk]['close'][-13:].mean()
        stock_poor[stk]['ma13'] = (data[stk]['close'][-13:-1].sum() + ticks[stk]['lastPrice'])/13
        stock_poor[stk]['ma13-1'] = data[stk]['close'][-14:-1].mean()
        # stock_poor[stk]['ma20'] = data[stk]['close'][-20:].mean()
        stock_poor[stk]['ma20'] = (data[stk]['close'][-20:-1].sum() + ticks[stk]['lastPrice'])/20
        stock_poor[stk]['ma20-1'] = data[stk]['close'][-21:-1].mean()
        # stock_poor[stk]['ma60'] = data[stk]['close'][-60:].mean()
        stock_poor[stk]['ma60'] = (data[stk]['close'][-60:-1].sum() + ticks[stk]['lastPrice'])/60
        stock_poor[stk]['ma60-1'] = data[stk]['close'][-61:-1].mean()
        
        超短强度 = math.atan((stock_poor[stk]['ma5']/stock_poor[stk]['ma5-1']-1)*100)*57.3
        短期强度 = math.atan((stock_poor[stk]['ma13']/stock_poor[stk]['ma13-1']-1)*100)*57.3
        中期强度 = math.atan((stock_poor[stk]['ma20']/stock_poor[stk]['ma20-1']-1)*100)*57.3
        长期强度 = math.atan((stock_poor[stk]['ma60']/stock_poor[stk]['ma60-1']-1)*100)*57.3
        # stock_poor[stk]['强度评分'] = 超短强度/50 + 短期强度/40 + 中期强度/21 + 长期强度/10
        # print(stk, stock_poor[stk]['强度评分'])
        p_tmp[stk] = 超短强度/50 + 短期强度/40 + 中期强度/21 + 长期强度/10
    p_tmp_sorted = sorted(p_tmp.items(), key=lambda d: d[1], reverse=True)
    # print(f'排序后的ETF列表：{p_tmp_sorted}')
    stock_poor_sorted = []
    for i in range(10):
        tmp = str(p_tmp_sorted[i]).split(',')[0].replace('(', '').replace("'", "")
        stock_poor_sorted.append(tmp)
    print(f'---- ETF排序前十列表：{stock_poor_sorted}')
    return stock_poor_sorted

# 获取持仓信息
def get_holdings(accountid, datatype):
    '''
    arg:
        accountid: 账户id, 
        datatype:
            'FUTURE':期货
            'STOCK':股票
            ......
    return:
    {股票名:{'手数':int, '持仓成本':float, '浮动盈亏':float, '可用余额':int}}
    '''
    PositionInfo_dict = {}
    resultlist = get_trade_detail_data(accountid, datatype, 'POSITION')
    for obj in resultlist:
        PositionInfo_dict[obj.m_strInstrumentID + '.' + obj.m_strExchangeID] = {
            '手数':obj.m_nVolume/100,
            '持仓成本':obj.m_dOpenPrice,
            '浮动盈亏':obj.m_dFloatProfit,
            '可用余额':obj.m_nCanUseVolume,
            '成交日期':obj.m_strOpenDate
        }
    return PositionInfo_dict

# 写入缓存变量到plk文件
def write_pickle(data:list, etf_bak_file:str):
    with open(etf_bak_file, 'wb') as f:
        pickle.dump(data, f)

# 读写本地缓存变量到内存
def load_pickle(file:str):
    with open(file, 'rb') as f:
        plk_data = pickle.load(f)
    return plk_data

# 交易函数
def trade(ContextInfo):
    now = datetime.now()
    # stock_poor_sorted = []
    stock_poor_sorted = etf_sort(ContextInfo, etf_file_path)
    etf_top3 = stock_poor_sorted[:3]
    # write_pickle(etf_top3, etf_bak_file)
    etf_bak_list = load_pickle(etf_bak_file)
    # holdings = get_holdings(ContextInfo.accID, 'STOCK')
    # 买入ETF排序前三
    for i in range(3):
        ticks = ContextInfo.get_full_tick(etf_top3)
        holdings = get_holdings(ContextInfo.accID, 'STOCK')
        print(f'持仓：{holdings}')
        if etf_top3[i] in holdings.keys():
            continue
        stk_price = ticks[stock_poor_sorted[i]]['lastPrice'] * 买入下单溢价
        stk_amount = int((买入金额/stk_price)/100)*100
        passorder(23 , 1101, ContextInfo.accID, stock_poor_sorted[i], 11, stk_price, stk_amount, 2, ContextInfo)  #立即下单
        print(f"买入：{stock_poor_sorted[i]}, 买入价：{stk_price}, 交易数量：{stk_amount}")
        etf_bak_list.append(etf_top3[i])
    
    # 卖出跌出前三的ETF
    for stk in etf_bak_list:
        if stk not in etf_top3 and holdings[stk]['可用余额'] > 0:
            ticks = ContextInfo.get_full_tick(stk)
            stk_price = ticks[stock_poor_sorted[i]]['lastPrice'] * 卖出下单溢价
            holdings = get_holdings(ContextInfo.accID, 'STOCK')
            stk_amount = holdings[stk]['手数']
            # passorder(24, 1101, ContextInfo.accID, stk, 42, stk_price, stk_amount, '',1,'',ContextInfo)
            order_target_percent(stk, 0, 'fix', stk_price, ContextInfo)
            etf_bak_list.remove(stk)
    write_pickle(etf_bak_list, etf_bak_file)

def handlebar(ContextInfo):
    test(ContextInfo)
    # trade(ContextInfo)
    pass

def test(ContextInfo):
    etf = etf_sort(ContextInfo, etf_file_path)
    # etf = etf[:3]
    write_pickle(etf, etf_bak_file)
    f = load_pickle(etf_bak_file)
    for i in f:
        name = ContextInfo.get_stock_name(i)
        print(f'排名前十的ETF：{i} {name}')



